package proxy

import (
	"fmt"
	"log"
	"net/http"
	"strconv"
	"strings"

	"github.com/gorilla/websocket"
)

type WebSocketProxy struct {
	outBoundConn *websocket.Conn
	proxyConn *websocket.Conn
	closeCh chan struct{}
}

func NewWebSocketProxy(outBoundConn *websocket.Conn, targetServer string, r *http.Request) (*WebSocketProxy, error) {
	var host string
	var port int
	if len(targetServer) == 0 || !strings.Contains(targetServer, ":")  {
		host = "localhost"
		port = 8080
	} else {
		arr := strings.Split(targetServer, ":")
		host = arr[0]
		port, _ = strconv.Atoi(arr[1])
	}

	uri := r.URL.RequestURI()
	proxyUrl := fmt.Sprintf("ws://%s:%d%s", host, port, uri)
	proxyConn, dailResponse, err := websocket.DefaultDialer.Dial(proxyUrl, http.Header{})
	if err != nil {
		log.Printf("dail error, err=%v, response=%v", err, dailResponse)
		return nil, err
	}

	proxy := &WebSocketProxy{
		outBoundConn: outBoundConn,
		proxyConn: proxyConn,
		closeCh: make(chan struct{}, 1),
	}
	return proxy, nil
}

func (w *WebSocketProxy) SendToProxy(messageType int, data []byte) error {
	err := w.proxyConn.WriteMessage(messageType, data)
	if err != nil {
		fmt.Printf("proxy websocket WriteMessage error, %v", err)
	}
	return err
}

func (w *WebSocketProxy) Start() {
LOOP: 
	for {
		select {
			case <- w.closeCh:
				break LOOP
			default:
				messsageType, data, err := w.proxyConn.ReadMessage()
				if err != nil {
					log.Printf("read websocket error, err=%v", err)
					break LOOP
				}

				err = w.outBoundConn.WriteMessage(messsageType, data)
				if err != nil {
					fmt.Printf("outBoundConn websocket WriteMessage error, %v", err)
				}
		}
	}
}

func (w *WebSocketProxy) Stop() {
	close(w.closeCh)
	w.proxyConn.Close()
}